from kittentts import KittenTTS
import os
import sys
import logging

# 确保可以导入本地模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 如果从包内导入失败，则尝试直接导入
try:
    from .config import LOG_LEVEL, TTS_MODEL_NAME, DEFAULT_LANGUAGE
except ImportError:
    from config import LOG_LEVEL, TTS_MODEL_NAME
    # 如果配置文件中没有DEFAULT_LANGUAGE，则设置默认值
    DEFAULT_LANGUAGE = "en-us"

# 配置日志
logging.basicConfig(level=getattr(logging, LOG_LEVEL))
logger = logging.getLogger(__name__)

class TTSProcessor:
    """KittenTTS处理器，封装模型调用逻辑"""
    def __init__(self):
        """初始化TTS处理器，加载KittenTTS模型"""
        try:
            # 初始化不同语言的模型字典
            self.models = {}
            # 默认加载英文模型
            self.models["en-us"] = KittenTTS(TTS_MODEL_NAME)
            # 加载中文模型
            self.models["zh"] = KittenTTS(TTS_MODEL_NAME, language="zh")
            
            os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
            logger.info(f"KittenTTS models loaded successfully for languages: {list(self.models.keys())}")
        except Exception as e:
            logger.error(f"Failed to load KittenTTS model: {e}")
            raise
    
    def generate_audio(self, text, voice='expr-voice-5-m', speed=1.0, language=DEFAULT_LANGUAGE):
        """生成音频数据
        
        Args:
            text: 要转换的文本
            voice: 使用的语音
            speed: 语速
            language: 使用的语言，默认为配置文件中的DEFAULT_LANGUAGE
            
        Returns:
            numpy数组形式的音频数据
        """
        try:
            # 确保语言参数有效
            if language not in self.models:
                # 如果请求的语言不存在，使用默认语言
                logger.warning(f"Language '{language}' not available, using default: {DEFAULT_LANGUAGE}")
                language = DEFAULT_LANGUAGE
                
            logger.info(f"Generating audio for text: {text[:30]}... with voice {voice}, speed {speed}, language {language}")
            audio = self.models[language].generate(text, voice=voice, speed=speed)
            logger.info(f"Audio generated successfully, shape: {audio.shape}")
            return audio
        except Exception as e:
            logger.error(f"Failed to generate audio: {e}")
            raise
    
    def get_available_voices(self):
        """获取可用的语音列表
        
        Returns:
            可用语音的列表
        """
        try:
            # 使用默认语言的模型获取语音列表
            return self.models[DEFAULT_LANGUAGE].available_voices
        except Exception as e:
            logger.error(f"Failed to get available voices: {e}")
            # 返回默认的语音列表作为后备
            return [
                'expr-voice-2-m', 'expr-voice-2-f', 
                'expr-voice-3-m', 'expr-voice-3-f', 
                'expr-voice-4-m', 'expr-voice-4-f', 
                'expr-voice-5-m', 'expr-voice-5-f'
            ]
    
    def get_available_languages(self):
        """获取可用的语言列表
        
        Returns:
            可用语言的列表
        """
        return list(self.models.keys())